// Copyright (c) 2024-present AI-Labs

// @ ts-nocheck
import { Extension } from "@tiptap/core";
import Suggestion, { SuggestionOptions, SuggestionProps } from "@tiptap/suggestion";

import tippy, { Instance } from "tippy.js";
import OpenAI from "openai";
import { sleep } from "openai/core.mjs";
import { t } from "i18next";

import { AICommand, markdownToHtml, uuid } from "@uai-team/uai-editor-common";
import { InnerEditor } from "@/core/UAIEditor.ts";

export type QuickCommandOptions = {
    HTMLAttributes?: Record<string, any>;
    suggestion: Omit<SuggestionOptions, 'editor'>;
}

/**
 * 定义快捷命令
 */
export default Extension.create<QuickCommandOptions>({
    name: 'quickCommand',
    addOptions() {
        return {
            suggestion: {
                char: '/',
                command: async ({ editor, range, props }) => {
                    editor.chain().focus().deleteRange(range).run();
                    const from = editor.state.selection.to;
                    const type = props.type;
                    const command = props.command as AICommand;
                    const options = (editor as InnerEditor).uaiEditor.options;
                    const actionKey = "uai_waiting_action";

                    const selectedText = editor.state.selection.$head.parent.textContent;
                    if (type === "chat") {
                        // 模型对话命令
                        var models = Object.keys(options.ai?.chat?.models ?? {});
                        if (models.length > 0) {
                            var model = command.model ?? "auto";
                            if (model === "auto") {
                                model = models[0];
                            }
                            const modelConfig = options.ai!.chat!.models![model];
                            const client = new OpenAI({
                                baseURL: modelConfig.baseUrl,
                                apiKey: modelConfig.apiKey,
                                dangerouslyAllowBrowser: true,
                            });

                            client.chat.completions.create({
                                model: modelConfig.model ?? "o1",
                                stream: true,
                                max_tokens: modelConfig.max_tokens,
                                temperature: modelConfig.temperature,
                                top_p: modelConfig.top_p,
                                frequency_penalty: modelConfig.frequency_penalty,
                                messages: [
                                    { "role": "system", "content": command.prompt ?? t('command.chat.prompt') },
                                    { "role": "user", "content": selectedText }
                                ],
                            }).then(async response => {
                                for await (var chunk of response) {
                                    var content = chunk.choices[0]?.delta?.content || '';
                                    editor.view.dispatch(editor.state.tr.insertText(content));
                                }
                                const end = editor.state.selection.to;
                                const insertedText = editor.state.doc.textBetween(from, end);
                                editor.view.dispatch(editor.state.tr.replaceWith(from, end, (editor as InnerEditor).parseHtml(markdownToHtml(insertedText))).scrollIntoView());
                            });
                        }
                    } else if (type === "text2image") {
                        // 文生图命令
                        var models = Object.keys(options.ai?.image?.models?.text2image ?? {});
                        if (models.length > 0) {
                            const id = uuid();
                            const previewType = "image";
                            const type = "image";
                            var model = command.model ?? "auto";
                            if (model === "auto") {
                                model = models[0];
                            }
                            const modelConfig = options.ai!.image!.models!.text2image![model];
                            const client = new OpenAI({
                                baseURL: modelConfig.baseUrl,
                                apiKey: modelConfig.apiKey,
                                dangerouslyAllowBrowser: true,
                            });
                            editor.view.dispatch(editor.state.tr.setMeta(actionKey, {
                                type: "add",
                                id,
                                pos: editor.state.tr.selection.from,
                            }));
                            client.images.generate({
                                "model": modelConfig.model ?? "dall-e-3",
                                "prompt": `${command.prompt ?? ""}${selectedText}`,
                                "size": "1024x1024",
                            }).then(response => {
                                editor.view.dispatch(editor.state.tr.setMeta(actionKey, { type: "remove", id }));
                                var src;

                                if (response.data) {
                                    var data = response.data[0];
                                    var url = response.data[0].url;
                                    var b64_json = response.data[0].b64_json;
                                    if (url) {
                                        src = url;
                                    } else if (b64_json && b64_json.startsWith("data:")) {
                                        src = b64_json;
                                    } else if (b64_json && !b64_json.startsWith("data:")) {
                                        src = `data:image;base64,${b64_json}`;
                                    } else {
                                        src = data;
                                    }
                                    // @ts-ignore
                                } else if (response.images) {
                                    // @ts-ignore
                                    var image = response.images[0];
                                    // @ts-ignore
                                    var url = response.images[0].url;
                                    // @ts-ignore
                                    var b64_json = response.images[0].b64_json;
                                    if (url) {
                                        src = url;
                                    } else if (b64_json && b64_json.startsWith("data:")) {
                                        src = b64_json;
                                    } else if (b64_json && !b64_json.startsWith("data:")) {
                                        src = `data:image;base64,${b64_json}`;
                                    } else {
                                        src = image;
                                    }
                                }

                                editor.commands.insertContentAt(editor.state.tr.selection.from, {
                                    type: type,
                                    attrs: {
                                        ['src']: src,
                                        type,
                                        previewType,
                                    },
                                });
                            });
                        }
                    } else if (type === "audio") {
                        var models = Object.keys(options.ai?.audio?.models?.speech ?? {});
                        if (models.length > 0) {
                            const id = uuid()
                            const previewType = "audio";
                            const type = "audio";
                            var model = command.model ?? "auto";
                            if (model === "auto") {
                                model = models[0];
                            }
                            const modelConfig = options.ai!.audio!.models!.speech![model];
                            editor.view.dispatch(editor.state.tr.setMeta(actionKey, {
                                type: "add",
                                id,
                                pos: editor.state.tr.selection.from,
                            }));
                            if (modelConfig.modelType === "moark" && modelConfig.model === "Spark-TTS-0.5B") {
                                await fetch(modelConfig.baseUrl!, {
                                    headers: modelConfig.headers,
                                    method: "POST",
                                    body: JSON.stringify({
                                        model: modelConfig.model ?? "tts-1",
                                        inputs: selectedText,
                                        gender: "female",
                                        pitch: 3,
                                        speed: 3,
                                    })
                                }).then(response => response.json())
                                    .then(async json => {
                                        let status = 'in_progress';
                                        let response: any;
                                        while (status != 'success') {
                                            await fetch(json.urls.get,
                                                {
                                                    headers: modelConfig.headers,
                                                    method: "GET",
                                                }
                                            ).then(response => response.json())
                                                .then(json => {
                                                    status = json.status;
                                                    response = json;
                                                });
                                            await sleep(5000);
                                        }
                                        editor.view.dispatch(editor.state.tr.setMeta(actionKey, { type: "remove", id }));
                                        editor.commands.insertContentAt(editor.state.tr.selection.from, {
                                            type: type,
                                            attrs: {
                                                ['src']: response.output.file_url,
                                                type,
                                                previewType,
                                            },
                                        });
                                    });

                            } else if (modelConfig.modelType === "moark") {
                                const response = await fetch(modelConfig.baseUrl!,
                                    {
                                        headers: modelConfig.headers,
                                        method: "POST",
                                        body: JSON.stringify({
                                            model: modelConfig.model ?? "tts-1",
                                            input: selectedText,
                                            voice: modelConfig.voice ?? "echo",
                                        })
                                    }
                                );
                                const blob = await response.blob();
                                const reader = new FileReader();
                                reader.readAsDataURL(blob);
                                reader.onloadend = () => {
                                    const base64data = reader.result as string;
                                    editor.view.dispatch(editor.state.tr.setMeta(actionKey, { type: "remove", id }));
                                    editor.commands.insertContentAt(editor.state.tr.selection.from, {
                                        type: type,
                                        attrs: {
                                            ['src']: base64data,
                                            type,
                                            previewType,
                                        },
                                    });
                                };
                            } else {
                                const client = new OpenAI({
                                    baseURL: modelConfig.baseUrl,
                                    apiKey: modelConfig.apiKey,
                                    dangerouslyAllowBrowser: true,
                                    defaultHeaders: modelConfig.headers
                                });
                                client.audio.speech.create({
                                    model: modelConfig.model ?? "tts-1",
                                    input: selectedText,
                                    voice: modelConfig.voice ?? "echo"
                                })
                                    .then(response => response.json())
                                    .then(json => {
                                        editor.view.dispatch(editor.state.tr.setMeta(actionKey, { type: "remove", id }));

                                        editor.commands.insertContentAt(editor.state.tr.selection.from, {
                                            type: type,
                                            attrs: {
                                                ['src']: json.data?.base64 ?? json.data?.url,
                                                type,
                                                previewType,
                                            },
                                        });
                                    });
                            }
                        }
                    } else if (type === "video") {
                        var models = Object.keys(options.ai?.video?.models ?? {});
                        const id = uuid()
                        if (models.length > 0) {
                            var model = command.model ?? "auto";
                            if (model === "auto") {
                                model = models[0];
                            }
                            const modelConfig = options.ai!.video!.models![model];
                            editor.view.dispatch(editor.state.tr.setMeta(actionKey, {
                                type: "add",
                                id,
                                pos: editor.state.tr.selection.from,
                            }));
                            fetch(modelConfig.baseUrl!, {
                                method: 'POST',
                                headers: modelConfig.headers,
                                body: JSON.stringify({
                                    api_key: modelConfig.apiKey,
                                    model: modelConfig.model,
                                    prompt: selectedText,
                                })
                            })
                                .then(response => response.json())
                                .then(data => {
                                    if (modelConfig.modelType === "CogVideoX") {
                                        var query = setInterval(() => {
                                            fetch(`https://open.bigmodel.cn/api/paas/v4/async-result/${data.id}`, {
                                                method: 'GET',
                                                headers: modelConfig.headers,
                                            })
                                                .then(response => response.json())
                                                .then(data => {
                                                    if (data.task_status !== "PROCESSING") {
                                                        clearInterval(query);
                                                        const previewType = "video";
                                                        const type = "video";
                                                        editor.view.dispatch(editor.state.tr.setMeta(actionKey, { type: "remove", id }));
                                                        if (data.task_status === "SUCCESS") {
                                                            editor.commands.insertContentAt(editor.state.tr.selection.from, {
                                                                type: type,
                                                                attrs: {
                                                                    ['src']: data.video_result[0].url, // cover_image_url
                                                                    type,
                                                                    previewType,
                                                                },
                                                            });
                                                        }
                                                    }
                                                })
                                        }, 5000);
                                    }
                                });
                        }
                    }
                },
                render: () => {
                    let container: HTMLElement;
                    let popup: Instance;

                    return {
                        onStart: (props: SuggestionProps) => {
                            container = document.createElement("div");
                            container.classList.add("uai-popup-action-list");
                            (props.editor as InnerEditor).uaiEditor.options.ai?.chat?.commands?.forEach(command => {
                                const item = document.createElement("div");
                                item.classList.add("uai-popup-action-item");
                                item.innerHTML = `${command.icon}&nbsp;&nbsp;${command.name}`;
                                item.addEventListener('click', () => {
                                    props.command({
                                        type: "chat",
                                        command: command
                                    })
                                });
                                container.appendChild(item);
                            })

                            container.appendChild(document.createElement("hr"));

                            (props.editor as InnerEditor).uaiEditor.options.ai?.image?.commands?.forEach(command => {
                                const item = document.createElement("div");
                                item.classList.add("uai-popup-action-item");
                                item.innerHTML = `${command.icon}&nbsp;&nbsp;${command.name}`;
                                item.addEventListener('click', () => {
                                    props.command({
                                        type: "text2image",
                                        command: command
                                    })
                                });
                                container.appendChild(item);
                            })

                            container.appendChild(document.createElement("hr"));

                            (props.editor as InnerEditor).uaiEditor.options.ai?.audio?.commands?.forEach(command => {
                                const item = document.createElement("div");
                                item.classList.add("uai-popup-action-item");
                                item.innerHTML = `${command.icon}&nbsp;&nbsp;${command.name}`;
                                item.addEventListener('click', () => {
                                    props.command({
                                        type: "audio",
                                        command: command
                                    })
                                });
                                container.appendChild(item);
                            })

                            container.appendChild(document.createElement("hr"));

                            (props.editor as InnerEditor).uaiEditor.options.ai?.video?.commands?.forEach(command => {
                                const item = document.createElement("div");
                                item.classList.add("uai-popup-action-item");
                                item.innerHTML = `${command.icon}&nbsp;&nbsp;${command.name}`;
                                item.addEventListener('click', () => {
                                    props.command({
                                        type: "video",
                                        command: command
                                    })
                                });
                                container.appendChild(item);
                            })

                            // @ts-ignore
                            popup = tippy('body', {
                                appendTo: props.editor.options.element,
                                getReferenceClientRect: props.clientRect,
                                content: container,
                                showOnCreate: true,
                                interactive: true,
                                allowHTML: true,
                                trigger: 'manual',
                                placement: 'right',
                                arrow: false,
                            })[0]
                        },
                        onUpdate(props) {
                            if (!props.clientRect) {
                                return;
                            }
                            popup.setProps({
                                getReferenceClientRect: props.clientRect as any,
                            })
                        },
                        onKeyDown(props) {
                            if (props.event.key === 'Escape') {
                                popup.hide();
                                return true;
                            }
                            return false;
                        },
                        onExit() {
                            popup.hide();
                            container.remove();
                        },
                    }
                },
            }
        }
    },
    addProseMirrorPlugins() {
        return [
            Suggestion({
                editor: this.editor,
                ...this.options.suggestion,
            }),
        ]
    },
})